#!/usr/bin/env python
# -*- coding: UTF-8 -*-

import sys, os, time
import logging
import collections
import types
from subprocess import Popen, PIPE
import traceback
import dbtool

SNMP_ROOT_OID = ".1.3.6.1.4.1.6101.3005"
SNMP_QUEYR_PARENT_OID = SNMP_ROOT_OID+'.1'
SNMP_QUEYR_PRODUCT_OID = SNMP_QUEYR_PARENT_OID+'.1'
SNMP_QUEYR_PATTERN_OID= SNMP_QUEYR_PARENT_OID+'.2.1'
SNMP_QUEYR_NIC_OID = SNMP_QUEYR_PARENT_OID+'.3.1'

MAX_PATTERN_CNT = 9
NIC_NAME_PREFIX='eth'


logging.basicConfig(filename='/opt/TrendMicro/DTAS/ManagementServer/bin/a.log',
                    level=logging.DEBUG,
                    format='%(levelname)s\t%(message)s [%(filename)s:%(lineno)d]')
logger = logging.getLogger("")

class SNMPException(Exception):
    pass

class SNMPInfo:
    '''
    The SNMP get vale common class
    '''
    def __init__(self,match_oid= None):
        self.match_oid = match_oid
        self.query_oid = None

        # SNMP print value and type
        self.value = None
        self.type  = 'string'

    def get_index_oid(self):
        '''
        Get all index from the oid
        Return type: []
        Maybe: 1 index,or 2 index
        '''
        logger.debug('Process OID [%s]'%self.query_oid)
        if self.query_oid is None:
            raise SNMPException('Child class must override the query_oid')

        if self.match_oid is None:
            raise SNMPException('Child class must override the match_oid')

        if self.query_oid.startswith(self.match_oid):
            return [int(i) for i in self.query_oid[len(self.match_oid):].split('.')[1:-1]]

        logger.error('query_oid [%s] dose not start with match_oid[%s]'%(self.query_oid,self.match_oid))
        return None

    def set_value_string(self,value=None):
        '''
        The SNMP query response value is a string
        '''
        self.value = value
        self.type = 'string'

    def set_value_integer(self,value=None):
        '''
        The SNMP query response value is an integer
        '''
        self.value = value
        self.type = 'integer'

    def get_value(self):
        '''
        Child class must override the method to set value and type for specific query OID
        '''
        pass

    def print_query_response(self):
        '''
        Print SNMP Query response
        '''

        if self.query_oid is None:
            raise SNMPException('Child class must override the query_oid')
        if self.value is None:
            raise SNMPException('Child class must set the value')
        logger.debug('Resp[%s]=%s'%(self.query_oid,str(self.value)))

        print self.query_oid
        print self.type
        print self.value

class ProductInfo(SNMPInfo):
    '''
    Not a table MIB,index will be one
    '''
    def __init__(self,query_oid):
        SNMPInfo.__init__(self,SNMP_QUEYR_PRODUCT_OID)
        self.query_oid = query_oid

    def __getenv(self,key):
        return os.getenv(key)

    def get_value(self):
        '''
        index_oid: The index of the object
        '''
        index_list = self.get_index_oid()

        if index_list is None:
            raise SNMPException('The get_index_oid return a None value')
        index_oid = index_list[0]
        value = ''
        if index_oid == 1 :
            value = self.__getenv('PRODUCT_VERSION')
        if index_oid == 2 :
            value = self.__getenv('PRODUCT_BUILD_NUM')
        if index_oid == 3 :
            value = self.__getenv('PRODUCT_HOTFIX_VERSION')
        logger.debug('Index %s value: %s'%(index_oid,value))

        return self.set_value_string(value)

class PatternInfo(SNMPInfo):
    '''
     A table MIB,index will be two
    '''
    def __init__(self,query_oid):
        SNMPInfo.__init__(self,SNMP_QUEYR_PATTERN_OID)

        # oid from the snmp manager/snmpwalk
        self.query_oid = query_oid

    def get_db(self):
        self.db_conn = dbtool.DBConnection('dtasdb', 'dtasuser', '', 'SNMP_MODE')
        try:
            self.db_conn.execute("select '0x'||to_hex(auid) auid,productname,downloadversion from appproduct T WHERE id > 11 ORDER BY id;")
            result = self.db_conn.fetchall()
            if len(result) == 0:
                logger.error('Could not get pattern/engine info from appproduct')
                return None
            else:
                return result
        except:
            logger.error(traceback.format_exc())
            return None

    def get_value(self):
        '''
        a pattern = (index,auid,name,version)
        :param query_oid:
        :return:
        '''
        index_list = self.get_index_oid()
        if index_list is None:
            raise SNMPException('The get_index_oid return a None value')
        which_info,which_index = index_list[0],index_list[1]
        list = self.get_db()
        if which_info == 1:
            return self.set_value_integer(which_index)

        logger.debug('list[%d][%d]='%(which_index-1,which_info-1,))
        value = list[which_index -1][which_info -2]
        logger.debug('=%s'%(value,))

        return self.set_value_string(value)

class ThroughtputInfo(SNMPInfo):
    '''
    A table MIB,index will be two
    '''
    def __init__(self,query_oid):
        SNMPInfo.__init__(self,SNMP_QUEYR_NIC_OID)
        self.query_oid = query_oid

    @classmethod
    def get_nic_cnt(cls):
        return int(cls._run_shell_cmd_with_output('cat /proc/net/dev|grep -c %s'%NIC_NAME_PREFIX))

    @classmethod
    def _run_shell_cmd_with_output(cls,cmd):
        try:
            handle = Popen(cmd, shell=True, stdout=PIPE)
            result = handle.communicate()[0].strip()
        except Exception, e:
            print e
        return result

    def _accumulate_throughput(self,dev_name, type):
        cmd = "cat /proc/net/dev | grep " + dev_name.strip()
        rst_str = self._run_shell_cmd_with_output(cmd).strip()
        rst_arr = rst_str[len(dev_name.strip())+1:].split()
        first_value = rst_arr[0] if type==1 else rst_arr[8]
        time.sleep(0.5)
        rst_str = self._run_shell_cmd_with_output(cmd).strip()
        rst_arr = rst_str[len(dev_name.strip())+1:].split()
        second_value = rst_arr[0] if type==1 else rst_arr[8]
        throughput = (float(second_value) - float(first_value))/0.5
        rst_value = "0 bps"
        if(throughput <= 0):
            return rst_value;
        elif(throughput < 1024):
            rst_value = "{0:.2f} bps".format(throughput)
        elif(throughput < 1024*1024):
            rst_value = "{0:.2f} Kbps".format(throughput/1024)
        elif(throughput < 1024*1024*1024):
            rst_value = "{0:.2f} Mbps".format(throughput/1024/1024)
        else:
            rst_value = "{0:.2f} Gbps".format(throughput/1024/1024/1024)
        return rst_value

    def get_value(self):
        '''
        index_oid: The index of the object
        '''
        index_list = self.get_index_oid()
        if index_list is None:
            raise SNMPException('The get_index_oid return a None value')
        which_info,which_index = index_list[0],index_list[1]
        if which_info == 1:
            return self.set_value_integer(which_index)
        nic_name = '%s%s'%(NIC_NAME_PREFIX,which_index-1)
        if which_info == 2:
            return self.set_value_string(nic_name)
        if which_info == 3:#Receive Throughput
            return  self.set_value_string(self._accumulate_throughput(nic_name,1))
        #Transmit Throughput
        return  self.set_value_string(self._accumulate_throughput(nic_name,2))


rootOIDmap=collections.OrderedDict([
    (SNMP_QUEYR_PRODUCT_OID+'.1.0', ProductInfo),
    (SNMP_QUEYR_PRODUCT_OID+'.2.0', ProductInfo),
    (SNMP_QUEYR_PRODUCT_OID+'.3.0', ProductInfo),
])

GET='-g'
GETNEXT='-n'



def snmp_walk(oidMap,action,QueryOid):

    keys = oidMap.keys()

    if action == GETNEXT:
        i = 0
        for k in keys:
            if k == QueryOid and i+1 < len(keys):
                oid = keys[i+1]
                if isinstance(oidMap[oid],types.ClassType) and issubclass(oidMap[oid],SNMPInfo):
                    info = oidMap[oid](oid)
                    info.get_value()
                    info.print_query_response()
                else:
                    raise SNMPException('Its not an SNMPInfo')
                break
            # oid is not match any, but has a prefix with current oid
            elif len(k) > len(QueryOid) and str(k).startswith(QueryOid):
                if isinstance(oidMap[k],types.ClassType) and issubclass(oidMap[k],SNMPInfo):
                    info = oidMap[k](k)
                    info.get_value()
                    info.print_query_response()
                else:
                    logger.debug(issubclass(oidMap[k],SNMPInfo))
                    raise SNMPException('Its not an SNMPInfo %s'%type(oidMap[k]))
                break
            i= i+1;
            logger.debug('k= %s QueryOid=%s'%(k,QueryOid))
    elif action == GET:
        if QueryOid in oidMap:
            if isinstance(oidMap[QueryOid],types.ClassType) and issubclass(oidMap[QueryOid],SNMPInfo):
                info = oidMap[QueryOid](QueryOid)
                info.get_value()
                info.print_query_response()
            else:
                raise SNMPException('Its not an SNMPInfo')
    else:
        logger.error('Do Not support action %s'%(action))
        return -1


def addMIBSequenceOID(rootOIDmap,parentOID,indexCnt,infoCnt,obj):
    '''
    The count of Sequence OID defined in MIB file is up to the specific OS and Not a fixed value
    '''
    for which_info in xrange(infoCnt):
        for which_index in xrange(indexCnt):
            rootOIDmap['%s.%s.%s.0'%(parentOID,which_info+1,which_index+1)]=obj

if __name__ == '__main__':
    addMIBSequenceOID(rootOIDmap,SNMP_QUEYR_PATTERN_OID,MAX_PATTERN_CNT,4,PatternInfo)
    addMIBSequenceOID(rootOIDmap,SNMP_QUEYR_NIC_OID,ThroughtputInfo.get_nic_cnt(),4,ThroughtputInfo)
    if len(sys.argv) != 3:
        print '''
Usage: %s [OPTION] OID
     -g    GET
     -n    GETNEXT
OID list:'''%(sys.argv[0],)
        for k in rootOIDmap.keys():
            print k
        sys.exit(-1)
    action = sys.argv[1].strip()
    oid = sys.argv[2].strip()
    logger.debug('INPUT OID:[%s] action: [%s]'%(oid,action))
    try:
        snmp_walk(rootOIDmap,action,oid)
    except Exception as e:
        logger.error(traceback.format_exc())

